# 序列化
from typing import Optional

from cryptography.hazmat.primitives import serialization
# 填充方案
from cryptography.hazmat.primitives import padding
# RSA 非对称加密库
import cryptography.hazmat.primitives.asymmetric.rsa as rsa_lib
import cryptography.hazmat.primitives.asymmetric.padding as rsa_padding
# AES 对称加密库
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
# 其他
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.exceptions import InvalidSignature
import os


class EncryptionModule:
    """
    加解密模块
    """
    def __init__(self):
        self._symmetric_key: Optional[algorithms.AES] = None
        self._symmetric_cipher: Optional[Cipher] = None
        self._asymmetric_local_private_key: Optional[rsa_lib.RSAPrivateKey] = None
        self._asymmetric_local_public_key: Optional[rsa_lib.RSAPublicKey] = None
        self._asymmetric_remote_public_key: Optional[rsa_lib.RSAPublicKey] = None

    def init_symmetric(self, key: bytes, iv: bytes) -> None:
        """
        初始化对称操作（AES 256 CBC）
        :param key: 对称密钥，32字节
        :param iv: 初始向量，16字节
        :return: 无
        """
        if self._symmetric_key is not None:
            del self._symmetric_key
        if self._symmetric_cipher is not None:
            del self._symmetric_cipher
        self._symmetric_key = algorithms.AES(key)
        self._symmetric_cipher = Cipher(self._symmetric_key, modes.CBC(iv))

    def encrypt_symmetric(self, plain_text: bytes):
        """
        对称加密，需要先调用init_symmetric
        :param plain_text: 明文
        :return: 密文
        """
        if self._symmetric_cipher is not None:
            # 填充
            padder = padding.PKCS7(algorithms.AES.block_size).padder()
            padded_data = padder.update(plain_text) + padder.finalize()

            # 加密
            encryptor = self._symmetric_cipher.encryptor()
            return encryptor.update(padded_data) + encryptor.finalize()
        else:
            return None

    def decrypt_symmetric(self, cipher_text: bytes):
        """
        对称解密，需要先调用init_symmetric
        :param cipher_text: 密文
        :return: 明文
        """
        if self._symmetric_cipher is not None:
            # 解密
            decryptor = self._symmetric_cipher.decryptor()
            decrypted_padded_data = decryptor.update(cipher_text) + decryptor.finalize()

            # 解填充
            padder = padding.PKCS7(algorithms.AES.block_size).unpadder()
            return padder.update(decrypted_padded_data) + padder.finalize()
        else:
            return None

    def generate_asymmetric_key(self) -> None:
        """
        生成非对称加密密钥，并自动导入模块
        :return: 无
        """
        if self._asymmetric_local_public_key is not None:
            del self._asymmetric_local_public_key
        if self._asymmetric_local_private_key is not None:
            del self._asymmetric_local_private_key
        self._asymmetric_local_private_key = rsa_lib.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
        self._asymmetric_local_public_key = self._asymmetric_local_private_key.public_key()

    def export_asymmetric_local_private_key(self) -> Optional[bytes]:
        """
        导出非对称本地私钥
        :return: pem格式的私钥数据
        """
        if self._asymmetric_local_private_key is not None:
            return self._asymmetric_local_private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
        else:
            return None

    def export_asymmetric_local_public_key(self) -> Optional[bytes]:
        """
        导出非对称本地公钥
        :return: pem格式的公钥数据
        """
        if self._asymmetric_local_public_key is not None:
            return self._asymmetric_local_public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            )
        else:
            return None

    def import_asymmetric_local_private_key(self, pem_path: str) -> None:
        """
        导入非对称本地私钥
        :param pem_path: pem密钥文件路径
        :return: 无
        """
        if not os.path.exists(pem_path):
            return None
        if self._asymmetric_local_public_key is not None:
            del self._asymmetric_local_public_key
        if self._asymmetric_local_private_key is not None:
            del self._asymmetric_local_private_key
        with open(pem_path, "rb") as key_file:
            self._asymmetric_local_private_key = serialization.load_pem_private_key(
                key_file.read(),
                password=None,
                backend=default_backend()
            )
            self._asymmetric_local_public_key = self._asymmetric_local_private_key.public_key()

    def import_asymmetric_remote_public_key(self, pem_path: str) -> None:
        """
        导入非对称远程公钥
        :param pem_path: pem密钥文件路径
        :return: 无
        """
        if not os.path.exists(pem_path):
            return None
        if self._asymmetric_remote_public_key is not None:
            del self._asymmetric_remote_public_key
        with open(pem_path, "rb") as f:
            self._asymmetric_remote_public_key = serialization.load_pem_public_key(
                f.read(),
                backend=default_backend()
            )

    def encrypt_asymmetric_local(self, plain_text: bytes) -> Optional[bytes]:
        """
        非对称加密，使用本地公钥加密，需要先导入本地私钥
        :param plain_text: 明文
        :return: 密文
        """
        if self._asymmetric_local_public_key is not None:
            left_len = len(plain_text)
            enc_offset = 0
            cipher_text = b''
            while left_len > 0:
                cur_enc_len = left_len
                if len(plain_text) > 128:
                    cur_enc_len = 128
                cipher_text += self._asymmetric_local_public_key.encrypt(
                    plain_text[enc_offset: enc_offset + cur_enc_len],
                    rsa_padding.OAEP(
                        mgf=rsa_padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None
                    )
                )
                enc_offset += cur_enc_len
                left_len -= cur_enc_len
            return cipher_text
        else:
            return None

    def encrypt_asymmetric_remote(self, plain_text: bytes) -> Optional[bytes]:
        """
        非对称加密，使用远端公钥加密，需要先导入远端公钥
        :param plain_text: 明文
        :return: 密文
        """
        if self._asymmetric_remote_public_key is not None:
            left_len = len(plain_text)
            enc_offset = 0
            cipher_text = b''
            while left_len > 0:
                cur_enc_len = left_len
                if left_len > 128:
                    cur_enc_len = 128
                cipher_text += self._asymmetric_remote_public_key.encrypt(
                    plain_text[enc_offset: enc_offset + cur_enc_len],
                    rsa_padding.OAEP(
                        mgf=rsa_padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None
                    )
                )
                enc_offset += cur_enc_len
                left_len -= cur_enc_len
            return cipher_text
        else:
            return None

    def decrypt_asymmetric(self, cipher_text: bytes) -> Optional[bytes]:
        """
        非对称解密，使用本地私钥解密，需要先导入本地私钥
        :param cipher_text: 密文
        :return: 明文
        """
        if self._asymmetric_local_private_key is not None:
            left_len = len(cipher_text)
            dec_offset = 0
            plain_text = b''
            while left_len > 0:
                cur_dec_len = left_len
                if left_len > 256:
                    cur_dec_len = 256
                plain_text += self._asymmetric_local_private_key.decrypt(
                    cipher_text[dec_offset: dec_offset + cur_dec_len],
                    rsa_padding.OAEP(
                        mgf=rsa_padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None
                    )
                )
                dec_offset += cur_dec_len
                left_len -= cur_dec_len
            return plain_text
        else:
            return None

    def sign_asymmetric(self, plain_text: bytes) -> Optional[bytes]:
        """
        签名，使用本地私钥，需要先导入本地私钥
        :param plain_text: 明文
        :return: 签名
        """
        if self._asymmetric_local_private_key is not None:
            return self._asymmetric_local_private_key.sign(
                plain_text,
                rsa_padding.PSS(
                    mgf=rsa_padding.MGF1(hashes.SHA256()),
                    salt_length=rsa_padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
        else:
            return None

    def verify_sign_asymmetric_local(self, plain_text: bytes, signature: bytes) -> Optional[bool]:
        """
        验证本地签名，需要先导入本地私钥
        :param plain_text: 明文
        :param signature: 签名
        :return: 验证结果
        """
        if self._asymmetric_local_public_key is not None:
            try:
                self._asymmetric_local_public_key.verify(
                    signature,
                    plain_text,
                    rsa_padding.PSS(
                        mgf=rsa_padding.MGF1(hashes.SHA256()),
                        salt_length=rsa_padding.PSS.MAX_LENGTH
                    ),
                    hashes.SHA256()
                )
            except InvalidSignature:
                return False
            else:
                return True
        else:
            return None

    def verify_sign_asymmetric_remote(self, plain_text: bytes, signature: bytes) -> Optional[bool]:
        """
        验证远端签名，需要先导入远端公钥
        :param plain_text: 明文
        :param signature: 签名
        :return: 验证结果
        """
        if self._asymmetric_remote_public_key is not None:
            try:
                self._asymmetric_remote_public_key.verify(
                    signature,
                    plain_text,
                    rsa_padding.PSS(
                        mgf=rsa_padding.MGF1(hashes.SHA256()),
                        salt_length=rsa_padding.PSS.MAX_LENGTH
                    ),
                    hashes.SHA256()
                )
            except InvalidSignature:
                return False
            else:
                return True
        else:
            return None
